package suno

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/go-redis/redis/v8"
	"github.com/imroc/req/v3"
	"gorm.io/gorm"
	"io"
	"myai/core/types"
	logger2 "myai/logger"
	"myai/service"
	"myai/service/oss"
	"myai/store"
	"myai/store/model"
	"myai/utils"
	"time"
)

var log = logger2.GetLogger()

type Service struct {
	httpClient  *req.Client
	taskQueue   *store.RedisQueue
	notifyQueue *store.RedisQueue
	db          *gorm.DB
	uploader    *oss.UploaderManager
	wsService   *service.WebSocketService
	userService *service.UserService
	clientIds   map[string]string
}

func NewService(db *gorm.DB, manager *oss.UploaderManager,
	redisCli *redis.Client, wsService *service.WebSocketService,
	userService *service.UserService) *Service {
	return &Service{
		httpClient:  req.C().SetTimeout(time.Minute * 3),
		taskQueue:   store.NewRedisQueue("My_Suno_Task_Queue", redisCli),
		notifyQueue: store.NewRedisQueue("My_Suno_Notify_Queue", redisCli),
		db:          db,
		wsService:   wsService,
		uploader:    manager,
		userService: userService,
		clientIds:   map[string]string{},
	}
}

type RespVo struct {
	Code    string `json:"code"`
	Message string `json:"message"`
	Data    string `json:"data"`
	Channel string `json:"channel,omitempty"`
}

type QueryRespVo struct {
	Code    string `json:"code"`
	Message string `json:"message"`
	Data    struct {
		TaskId     string `json:"task_id"`
		Action     string `json:"action"`
		Status     string `json:"status"`
		FailReason string `json:"fail_reason"`
		SubmitTime int    `json:"submit_time"`
		StartTime  int    `json:"start_time"`
		FinishTime int    `json:"finish_time"`
		Progress   string `json:"progress"`
		Data       []struct {
			Id       string `json:"id"`
			Title    string `json:"title"`
			Status   string `json:"status"`
			Metadata struct {
				Tags         string      `json:"tags"`
				Type         string      `json:"type"`
				Prompt       string      `json:"prompt"`
				Stream       bool        `json:"stream"`
				Duration     float64     `json:"duration"`
				ErrorMessage interface{} `json:"error_message"`
			} `json:"metadata"`
			AudioUrl          string `json:"audio_url"`
			ImageUrl          string `json:"image_url"`
			VideoUrl          string `json:"video_url"`
			ModelName         string `json:"model_name"`
			DisplayName       string `json:"display_name"`
			ImageLargeUrl     string `json:"image_large_url"`
			MajorModelVersion string `json:"major_model_version"`
		} `json:"data"`
	} `json:"data"`
}

func (s *Service) Run() {
	// 将数据库中未提交的任务加载到队列中
	var jobs []model.SunoJob
	s.db.Where("task_id", "").Where("progress", 0).Find(&jobs)
	for _, job := range jobs {
		var task types.SunoTask
		err := utils.JsonDecode(job.TaskInfo, &task)
		if err != nil {
			log.Errorf("Suno 解码任务信息时出错：%v", err)
			continue
		}
		task.Id = job.Id
		s.PushTask(task)
		s.clientIds[job.TaskId] = task.ClientId
	}
	log.Info("启动Suno任务消费者 ...")
	go func() {
		for {
			var task types.SunoTask
			err := s.taskQueue.LPop(&task)
			if err != nil {
				log.Errorf("Suno 消息任务队列任务出现错误：%v", err)
				continue
			}
			var r RespVo
			if task.Type == 3 && task.SongId != "" { // 歌曲拼接
				r, err = s.Merge(task)
			} else if task.Type == 4 && task.AudioURL != "" { // 上传歌曲
				r, err = s.Upload(task)
			} else { // 歌曲创作
				r, err = s.Create(task)
			}
			if err != nil {
				log.Errorf("Suno 任务执行失败：%v", err)
				s.db.Model(&model.SunoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
					"err_msg":  err.Error(),
					"progress": service.FailTaskProgress,
				})
				// 通知前端
				s.notifyQueue.RPush(service.NotifyMessage{
					ClientId: task.ClientId,
					UserId:   task.UserId,
					JobId:    int(task.Id),
					Message:  service.TaskStatusFailed,
				})
				continue
			}
			// 更新任务信息
			s.db.Model(&model.SunoJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
				"task_id": r.Data,
				"channel": r.Channel,
			})
			log.Infof("任务提交成功")
			s.clientIds[r.Data] = task.ClientId
			log.Infof("任务提交成功客户端信息：%v", s.clientIds)
		}
	}()
}

func (s *Service) PushTask(task types.SunoTask) {
	log.Infof("将新的 Suno 任务到添加到任务列表：%+v", task)
	s.taskQueue.RPush(task)
}

// Create 生成歌曲
func (s *Service) Create(task types.SunoTask) (RespVo, error) {
	// 读取 API KEY
	var apiKey model.ApiKey
	session := s.db.Session(&gorm.Session{}).Where("type", "suno").Where("enabled", true)
	if task.Channel != "" {
		session = session.Where("api_url", task.Channel)
	}
	tx := session.Order("last_used_at desc").First(&apiKey)
	if tx.Error != nil {
		return RespVo{}, errors.New("没有可用于Suno的API密钥")
	}
	reqBody := map[string]interface{}{
		"task_id":           task.RefTaskId,    // 续写的前任务id
		"continue_clip_id":  task.RefSongId,    // 需要继续创作时使用。含义为，需要继续创作的歌曲 id
		"continue_at":       task.ExtendSecs,   // 需要继续创作时使用。含义为，从第几秒开始继续创作，例如 120.00 或者 61.59
		"make_instrumental": task.Instrumental, // 是否纯音乐，默认不是纯音乐
		"mv":                task.Model,        // mv模型，chirp-v3-0、chirp-v3-5。不写默认 chirp-v3-0
	}
	// 灵感模式
	if task.Type == 1 {
		// 灵感模式至多200字(一个中文算1个字);传知名艺术家名称100%不会生成,比如周杰伦/郎朗
		reqBody["gpt_description_prompt"] = task.Prompt
	} else {
		// 自定义模式
		// prompt 歌词 (自定义模式专用)
		reqBody["prompt"] = task.Prompt
		// 风格标签(自定义模式专用)
		reqBody["tags"] = task.Tags
		// 标题(自定义模式专用)
		reqBody["title"] = task.Title
	}
	var res RespVo
	apiURL := fmt.Sprintf(service.SunoGenerateMusic, apiKey.ApiURL)
	log.Debugf("Suno：请求地址：%s，请求参数：%+v", apiURL, reqBody)
	r, err := s.httpClient.R().
		SetHeader("Authorization", "Bearer "+apiKey.Value).
		SetBody(reqBody).
		Post(apiURL)
	if err != nil {
		return RespVo{}, fmt.Errorf("Suno请求API：%s 出错：%v", apiURL, err)
	}
	body, _ := io.ReadAll(r.Body)
	err = json.Unmarshal(body, &res)
	if err != nil {
		return RespVo{}, fmt.Errorf("Suno解析API返回数据失败：%v, %s", err, string(body))
	}
	if res.Code != "success" {
		return RespVo{}, fmt.Errorf("请求Suno API 返回失败：%s", res.Message)
	}
	// 更新 API KEY 最后一次使用时间
	apiKey.LastUsedAt = time.Now().Unix()
	session.Updates(&apiKey)
	// 设置通道为，我们这次请求使用的API KEY 对应的 API 地址
	res.Channel = apiKey.ApiURL
	return res, nil
}

// Merge 歌曲拼接
func (s *Service) Merge(task types.SunoTask) (RespVo, error) {
	// 读取 API KEY
	var apiKey model.ApiKey
	session := s.db.Session(&gorm.Session{}).Where("type", "suno").Where("enabled", true)
	if task.Channel != "" {
		session = session.Where("api_url", task.Channel)
	}
	tx := session.Order("last_used_at desc").First(&apiKey)
	if tx.Error != nil {
		return RespVo{}, errors.New("没有可用于Suno的API密钥")
	}
	reqBody := map[string]interface{}{
		// extend 后的 歌曲ID
		"clip_id":   task.SongId,
		"is_infill": false,
	}

	var res RespVo
	apiURL := fmt.Sprintf(service.SunoConcatMusic, apiKey.ApiURL)
	log.Debugf("Suno：请求地址：%s，请求参数：%+v", apiURL, reqBody)
	r, err := s.httpClient.R().
		SetHeader("Authorization", "Bearer "+apiKey.Value).
		SetBody(reqBody).
		Post(apiURL)
	if err != nil {
		return RespVo{}, fmt.Errorf("Suno请求API：%s 出错：%v", apiURL, err)
	}
	body, _ := io.ReadAll(r.Body)
	err = json.Unmarshal(body, &res)
	if err != nil {
		return RespVo{}, fmt.Errorf("Suno解析API返回数据失败：%v, %s", err, string(body))
	}
	if res.Code != "success" {
		return RespVo{}, fmt.Errorf("请求Suno API 返回失败：%s", res.Message)
	}
	// 更新 API KEY 最后一次使用时间
	apiKey.LastUsedAt = time.Now().Unix()
	session.Updates(&apiKey)
	// 设置通道为，我们这次请求使用的API KEY 对应的 API 地址
	res.Channel = apiKey.ApiURL
	return res, nil
}

// Upload 上传参考音频
func (s *Service) Upload(task types.SunoTask) (RespVo, error) {
	// 读取 API KEY
	var apiKey model.ApiKey
	session := s.db.Session(&gorm.Session{}).Where("type", "suno").Where("enabled", true)
	if task.Channel != "" {
		session = session.Where("api_url", task.Channel)
	}
	tx := session.Order("last_used_at desc").First(&apiKey)
	if tx.Error != nil {
		return RespVo{}, errors.New("没有可用于Suno的API密钥")
	}
	reqBody := map[string]interface{}{
		// 参考音频地址
		"url": task.AudioURL,
	}

	var res RespVo
	apiURL := fmt.Sprintf(service.SunoUploadsAudioUrl, apiKey.ApiURL)
	log.Debugf("Suno：请求地址：%s，请求参数：%+v", apiURL, reqBody)
	r, err := s.httpClient.R().
		SetHeader("Authorization", "Bearer "+apiKey.Value).
		SetBody(reqBody).
		Post(apiURL)
	if err != nil {
		return RespVo{}, fmt.Errorf("Suno请求API：%s 出错：%v", apiURL, err)
	}
	if r.StatusCode != 200 {
		return RespVo{}, fmt.Errorf("Suno请求API：%s 状态码：%d，返回数据：%s", apiURL, r.StatusCode, r.String())
	}
	body, _ := io.ReadAll(r.Body)
	err = json.Unmarshal(body, &res)
	if err != nil {
		return RespVo{}, fmt.Errorf("Suno解析API返回数据失败：%v, %s", err, string(body))
	}
	if res.Code != "success" {
		return RespVo{}, fmt.Errorf("请求Suno API 返回失败：%s", res.Message)
	}
	// 更新 API KEY 最后一次使用时间
	apiKey.LastUsedAt = time.Now().Unix()
	session.Updates(&apiKey)
	// 设置通道为，我们这次请求使用的API KEY 对应的 API 地址
	res.Channel = apiKey.ApiURL
	return res, nil
}

// CheckTaskNotify 监听消息通知队列
func (s *Service) CheckTaskNotify() {
	go func() {
		log.Info("正在运行Suno任务通知检查 ...")
		for {
			var message service.NotifyMessage
			err := s.notifyQueue.LPop(&message)
			if err != nil {
				log.Errorf("Suno 从消息通知队列消费消息出现错误：%v", err)
				continue
			}
			log.Infof("Suno 消息通知：%+v", message)
			log.Infof("CheckTaskNotify客户端Id：%+v", s.clientIds)
			// 获取客户端连接
			client := s.wsService.Clients.Get(message.ClientId)
			if client == nil {
				log.Warnf("Suno 客户端连接为空：%+v", s.wsService.Clients)
				continue
			}
			// 发送消息
			utils.SendChannelMsg(client, types.ChSuno, message.Message)
		}
	}()
}

// DownloadFiles 下载文件
func (s *Service) DownloadFiles() {
	go func() {
		var items []model.SunoJob
		for {
			// 102 标识任务执行完成，但是文件未下载
			res := s.db.Where("progress", 102).Find(&items)
			if res.Error != nil {
				log.Errorf("Suno 下载文件时，查询数据库任务出错：%v", res.Error)
				time.Sleep(time.Second * 5)
				continue
			}
			for _, v := range items {
				// 下载图片和音频
				log.Infof("Suno 尝试下载封面图片：%s", v.CoverURL)
				coverURL, err := s.uploader.GetUploadHandler().PutUrlFile(v.CoverURL, true)
				if err != nil {
					log.Errorf("Suno 下载封面图片失败：%v", err)
					continue
				}
				log.Infof("Suno 尝试下载音频文件：%s", v.AudioURL)
				audioURL, err := s.uploader.GetUploadHandler().PutUrlFile(v.AudioURL, true)
				if err != nil {
					log.Errorf("Suno 下载音频文件失败：%v", err)
					continue
				}
				v.CoverURL = coverURL
				v.AudioURL = audioURL
				v.Progress = 100
				s.db.Updates(&v)
				// 通知前端
				s.notifyQueue.RPush(service.NotifyMessage{
					ClientId: s.clientIds[v.TaskId],
					UserId:   int(v.UserId),
					JobId:    int(v.Id),
					Message:  service.TaskStatusFinished,
				})
			}
			time.Sleep(time.Second * 10)
		}
	}()
}

// SyncTaskProgress 异步拉取任务，判断任务是否执行成功
func (s *Service) SyncTaskProgress() {
	go func() {
		var jobs []model.SunoJob
		for {
			// 查询进度小于100，并且任务Id不为空的任务
			res := s.db.Where("progress < ?", 100).Where("task_id <> ?", "").Find(&jobs)
			if res.Error != nil {
				log.Errorf("Suno SyncTaskProgress 查询数据库任务失败：%v", res.Error.Error())
				continue
			}
			for _, job := range jobs {
				// 查询任务
				task, err := s.QueryTask(job.TaskId, job.Channel)
				if err != nil {
					log.Errorf("Suno 查询任务失败：%v", err.Error())
					continue
				}

				if task.Code != "success" {
					log.Errorf("Suno 查询任务请求返回失败：%v", task)
					continue
				}

				log.Infof("Suno 任务返回数据：%+v", task)
				// 任务完成，删除旧的任务插入两条新任务
				if task.Data.Status == "SUCCESS" {
					var jobId = job.Id
					var flag = false
					// 开启事务
					tx := s.db.Begin()
					for _, v := range task.Data.Data {
						job.Id = 0
						job.Progress = 102 // 102 标识资源未下载完成
						job.Title = v.Title
						job.SongId = v.Id
						job.Duration = int(v.Metadata.Duration)
						job.Prompt = v.Metadata.Prompt
						job.Tags = v.Metadata.Tags
						job.ModelName = v.ModelName
						job.RawData = utils.JsonEncode(v)
						job.CoverURL = v.ImageLargeUrl
						job.AudioURL = v.AudioUrl

						if err = tx.Create(&job).Error; err != nil {
							log.Errorf("Suno 创新新任务失败：%v", err)
							tx.Rollback()
							break
						}
						flag = true
					}
					// 删除旧任务
					if flag {
						if err = tx.Delete(&model.SunoJob{}, "id = ?", jobId).Error; err != nil {
							log.Errorf("Suno 删除旧任务失败：%v", err)
							tx.Rollback()
							continue
						}
					}
					// 提交事务
					tx.Commit()
					// 通知前端
					log.Infof("客户端Id：%+v，任务Id：%s", s.clientIds, job.TaskId)
					s.notifyQueue.RPush(service.NotifyMessage{
						ClientId: s.clientIds[job.TaskId],
						UserId:   int(job.UserId),
						JobId:    int(job.Id),
						Message:  service.TaskStatusFinished,
					})
				} else if task.Data.FailReason != "" {
					// 任务执行失败
					job.Progress = service.FailTaskProgress
					job.ErrMsg = task.Data.FailReason
					s.db.Updates(&job)
					// 通知前端
					s.notifyQueue.RPush(service.NotifyMessage{
						ClientId: s.clientIds[job.TaskId],
						UserId:   int(job.UserId),
						JobId:    int(job.Id),
						Message:  service.TaskStatusFailed,
					})
				}
			}
			// 找出失败的任务，并恢复其扣减的算力
			s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
			for _, job := range jobs {
				err := s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
					Type:   types.PowerRefund,
					Model:  job.ModelName,
					Remark: fmt.Sprintf("Suno 任务失败，退回算力。任务Id：%s，Err：%s", job.TaskId, job.ErrMsg),
				})
				if err != nil {
					log.Errorf("Suno 任务失败，退回算力时出现错误：%v", err.Error())
					continue
				}
				// 更新任务状态
				s.db.Model(&job).UpdateColumn("power", 0)
			}
			time.Sleep(time.Second * 10)
		}
	}()
}

// QueryTask 查询任务
func (s *Service) QueryTask(taskId string, channel string) (QueryRespVo, error) {
	var apiKey model.ApiKey
	err := s.db.Session(&gorm.Session{}).Where("type", "suno").
		Where("api_url", channel).
		Where("enabled", true).
		Order("last_used_at desc").First(&apiKey).Error
	if err != nil {
		return QueryRespVo{}, errors.New("没有可用于Suno的API密钥")
	}
	apiURL := fmt.Sprintf(service.SunoFetch, apiKey.ApiURL, taskId)
	var res QueryRespVo
	// 发送请求
	r, err := s.httpClient.R().SetHeader("Authorization", "Bearer "+apiKey.Value).Get(apiURL)
	if err != nil {
		return QueryRespVo{}, fmt.Errorf("Suno请求API：%s 出错：%v", apiURL, err)
	}
	defer r.Body.Close()
	body, _ := io.ReadAll(r.Body)
	err = json.Unmarshal(body, &res)
	if err != nil {
		return QueryRespVo{}, fmt.Errorf("Suno解析API返回数据失败：%v, %s", err, string(body))
	}
	return res, nil
}
